package cn.funeralobjects.project.sshd.utils;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.FileUtils;
import org.apache.sshd.common.config.keys.BuiltinIdentities;
import org.apache.sshd.common.config.keys.KeyUtils;
import org.apache.sshd.common.keyprovider.KeyPairProvider;
import org.apache.sshd.common.keyprovider.MappedKeyPairProvider;
import org.apache.sshd.common.util.GenericUtils;
import org.apache.sshd.common.util.ValidateUtils;
import org.apache.sshd.common.util.io.resource.PathResource;
import org.apache.sshd.common.util.security.SecurityUtils;
import org.apache.sshd.server.keyprovider.AbstractGeneratorHostKeyProvider;
import org.apache.sshd.server.keyprovider.SimpleGeneratorHostKeyProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.Base64Utils;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.GeneralSecurityException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.PublicKey;
import java.security.spec.RSAPublicKeySpec;
import java.util.*;

/**
 * @author FuneralObjects Create date: 2020/8/5 4:43 AM
 */
public class PublicKeyUtils {

    private static final String RSA_KEY = "ssh-rsa";

    private static Logger log = LoggerFactory.getLogger(PublicKeyUtils.class);

    private static PublicKey decodePublicKey(byte[] keyBytes) throws GeneralSecurityException {
        // strip any newline chars
        while (keyBytes[keyBytes.length - 1] == 0x0a || keyBytes[keyBytes.length - 1] == 0x0d) {
            keyBytes = Arrays.copyOf(keyBytes, keyBytes.length - 1);
        }
        byte[] decodeBuffer = Base64Utils.decode(keyBytes);
        ByteBuffer bb = ByteBuffer.wrap(decodeBuffer);
        int len = bb.getInt();
        byte[] type = new byte[len];
        bb.get(type);
        if (RSA_KEY.equals(new String(type))) {
            BigInteger e = decodeBigInt(bb);
            BigInteger m = decodeBigInt(bb);
            RSAPublicKeySpec spec = new RSAPublicKeySpec(m, e);
            return KeyFactory.getInstance("RSA").generatePublic(spec);

        } else {
            throw new IllegalArgumentException("Only supports RSA");
        }
    }

    private static BigInteger decodeBigInt(ByteBuffer bb) {
        int len = bb.getInt();
        byte[] bytes = new byte[len];
        bb.get(bytes);
        return new BigInteger(bytes);
    }

    public static PublicKey readPubFile(File file) throws IOException {
        String pubContent = FileUtils.readFileToString(file, StandardCharsets.UTF_8);
        try {
            if (pubContent.startsWith(RSA_KEY)) {
                return decodePublicKey(pubContent.split(" ")[1].getBytes());
            }
            return decodePublicKey(pubContent.split(" ")[0].getBytes());
        } catch (GeneralSecurityException e) {
            throw new IOException(e);
        }

    }

    public static KeyPairProvider resolveServerKeys(String hostKeyType, int hostKeySize, Collection<String> keyFiles) throws Exception {
        if (GenericUtils.isEmpty(keyFiles)) {
            AbstractGeneratorHostKeyProvider hostKeyProvider;
            Path hostKeyFile;
            if (SecurityUtils.isBouncyCastleRegistered()) {
                hostKeyFile = Paths.get("key.pem");
                hostKeyProvider = SecurityUtils.createGeneratorHostKeyProvider(hostKeyFile);
            } else {
                hostKeyFile = Paths.get("key.ser");
                hostKeyProvider = new SimpleGeneratorHostKeyProvider(hostKeyFile);
            }
            hostKeyProvider.setAlgorithm(hostKeyType);
            if (hostKeySize != 0) {
                hostKeyProvider.setKeySize(hostKeySize);
            }
            List<KeyPair> keys = ValidateUtils.checkNotNullAndNotEmpty(hostKeyProvider.loadKeys(null), "Failed to load keys from %s", hostKeyFile);
            KeyPair kp = keys.get(0);
            PublicKey pubKey = kp.getPublic();
            String keyAlgorithm = pubKey.getAlgorithm();
            if (BuiltinIdentities.Constants.ECDSA.equalsIgnoreCase(keyAlgorithm)) {
                keyAlgorithm = KeyUtils.EC_ALGORITHM;
            } else if (BuiltinIdentities.Constants.ED25519.equals(keyAlgorithm)) {
                keyAlgorithm = SecurityUtils.EDDSA;
            }
            // force re-generation of host key if not same algorithm
            if (!Objects.equals(keyAlgorithm, hostKeyType)) {
                Files.deleteIfExists(hostKeyFile);
                hostKeyProvider.clearLoadedKeys();
            }
            return hostKeyProvider;
        } else {
            List<KeyPair> pairs = new ArrayList<>(keyFiles.size());
            for (String keyFilePath : keyFiles) {
                Path path = Paths.get(keyFilePath);
                PathResource location = new PathResource(path);
                Iterable<KeyPair> ids;
                try (InputStream inputStream = location.openInputStream()) {
                    ids = SecurityUtils.loadKeyPairIdentities(null, location, inputStream, null);
                } catch (Exception e) {
                    log.error("failed to load host key file {}", keyFilePath, e);
                    throw e;
                }

                if (ids == null) {
                    log.warn("no keys loaded from {}", keyFilePath);
                    continue;
                }

                for (KeyPair kp : ids) {
                    if (kp == null) {
                        log.warn("empty key found in {}", keyFilePath);
                        continue; // debug breakpoint
                    }
                    pairs.add(kp);
                }
            }
            return new MappedKeyPairProvider(
                    ValidateUtils.checkNotNullAndNotEmpty(pairs, "No key pairs loaded for provided key files"));
        }
    }

}
